#!/usr/bin/env python3
"""
Inscode Agent 命令行客户端
提供终端交互界面与AI助手进行对话

MCP支持需要安装：https://github.com/wonderwhy-er/DesktopCommanderMCP
注意：需要Node.js版本v20或以上
"""

import asyncio
import json
import sys
import os
from typing import Optional, List, Dict, Any, Union
import argparse
import logging
from dotenv import load_dotenv
from langchain_core.messages import AnyMessage

from inscode_agent import BaseGraph

from inscode_agent.graph.nodes.mcp import MCPNodes
from rich.console import Console
from rich.panel import Panel
from rich.markdown import Markdown
from rich.rule import Rule

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from inscode_agent.graph.coder import CoderGraph
from inscode_agent.mcps import MCPClientManager
from inscode_agent.models.graph import GraphRequest, GraphNodeParam, AgentStepLog

# 配置日志 - 完全禁用所有日志输出
logging.basicConfig(level=logging.CRITICAL)
logger = logging.getLogger(__name__)

# 禁用所有相关模块的日志输出
logging.getLogger("inscode_agent").setLevel(logging.CRITICAL)
logging.getLogger("langfuse").setLevel(logging.CRITICAL)
logging.getLogger("httpx").setLevel(logging.CRITICAL)
logging.getLogger("httpcore").setLevel(logging.CRITICAL)
logging.getLogger("urllib3").setLevel(logging.CRITICAL)
logging.getLogger("openai").setLevel(logging.CRITICAL)
logging.getLogger("asyncio").setLevel(logging.CRITICAL)

# 禁用loguru日志输出
try:
    from loguru import logger as loguru_logger

    loguru_logger.remove()
    loguru_logger.add(lambda _: None)  # 添加一个空的处理器
except ImportError:
    pass


# 重定向标准错误输出以捕获其他日志
class NullWriter:
    def write(self, txt):
        pass

    def flush(self):
        pass


# 保存原始stderr以便在需要时恢复
original_stderr = sys.stderr


# 全局对话历史缓存
class ConversationCache:
    """对话历史缓存管理器"""

    def __init__(self):
        from collections import OrderedDict
        self.conversations: OrderedDict[str, List[Dict[str, Any]]] = OrderedDict()

    def add_message(self, thread_id: str, message: Union[List[AnyMessage], AnyMessage]):
        """添加消息到指定线程"""
        if thread_id not in self.conversations:
            self.conversations[thread_id] = []

        # 处理消息可能是数组或单个消息的情况
        if isinstance(message, list):
            self.conversations[thread_id].extend(message)
        else:
            self.conversations[thread_id].append(message)

    def get_history(self, thread_id: str) -> List[AnyMessage]:
        """获取指定线程的历史记录"""
        return self.conversations.get(thread_id, [])

    def get_all_history(self) -> List[Dict[str, Any]]:
        """获取所有历史记录，按thread_id顺序"""
        all_history = []
        for thread_id in self.conversations:
            all_history.extend(self.conversations[thread_id])
        return all_history

    @property
    def thread_order(self) -> List[str]:
        """获取按顺序排列的thread_id列表"""
        return list(self.conversations.keys())

    def clear_thread(self, thread_id: str):
        """清除指定线程的历史记录"""
        if thread_id in self.conversations:
            del self.conversations[thread_id]

    def clear_all(self):
        """清除所有历史记录"""
        self.conversations.clear()


# 全局对话缓存实例
conversation_cache = ConversationCache()


class InsCodeClient:
    """Inscode Agent 命令行客户端"""

    def __init__(self):
        self.console = Console()
        self.graph: Optional[CoderGraph] = None
        self.thread_id = None
        self.mcp_client = None

    def check_required_files(self):
        """检查必需的文件是否存在"""
        required_files = ['.env', 'mcp.json']
        missing_files = []

        for file in required_files:
            if not os.path.exists(file):
                missing_files.append(file)

        if missing_files:
            self.console.print(f"[bold red]❌ 缺少必需文件[/bold red]")
            self.console.print(f"当前目录: [dim]{os.getcwd()}[/dim]")
            self.console.print("缺少以下文件:")
            for file in missing_files:
                self.console.print(f"  • [red]{file}[/red]")
            self.console.print("请确保在包含 [cyan].env[/cyan] 和 [cyan]mcp.json[/cyan] 文件的目录中运行此程序。")
            return False

        return True

    def load_environment(self):
        """加载环境变量"""
        try:
            # 临时禁用stderr输出
            sys.stderr = NullWriter()

            # 加载 .env 文件
            load_dotenv()

            # 恢复stderr
            sys.stderr = original_stderr

            self.console.print("[green]✅ 环境变量加载完成[/green]")
            return True
        except Exception as e:
            # 恢复stderr
            sys.stderr = original_stderr
            self.console.print(f"[red]❌ 加载环境变量失败: {str(e)}[/red]")
            return False

    def print_welcome(self):
        """打印欢迎信息"""
        self.console.print()
        self.console.print("[bold blue]🚀 Inscode AI Assistant[/bold blue]")
        self.console.print()
        self.console.print("欢迎使用 Inscode AI 助手！我是您的专业软件工程师AI助手，")
        self.console.print("拥有丰富的编程语言、框架、设计模式和最佳实践知识。")
        self.console.print()
        self.console.print("[bold]我可以帮助您：[/bold]")
        self.console.print("  • 编写和修改代码")
        self.console.print("  • 分析项目结构")
        self.console.print("  • 执行命令和脚本")
        self.console.print("  • 搜索和获取信息")
        self.console.print("  • 处理图像和文件")
        self.console.print()
        self.console.print("[dim]输入 [cyan]/help[/cyan] 查看帮助信息，输入 [cyan]/quit[/cyan] 退出程序[/dim]")
        self.console.print()

    async def initialize_graph(self):
        """初始化CoderGraph实例"""
        try:

            # 临时禁用stderr输出以避免初始化过程中的日志
            sys.stderr = NullWriter()

            # MCP配置文件路径（已确认存在）
            mcp_config_path = "mcp.json"

            try:
                # 初始化MCP客户端管理器
                self.mcp_client = MCPClientManager(mcp_config_path)
                await self.mcp_client.initialize()
                # 恢复stderr
                sys.stderr = original_stderr
                self.console.print("[green]✅ MCP客户端初始化完成[/green]")
            except Exception as mcp_error:
                # 恢复stderr
                sys.stderr = original_stderr
                self.console.print(f"[yellow]⚠️  MCP初始化失败: {mcp_error}[/yellow]")
                self.console.print("[yellow]将继续使用基础功能[/yellow]")
                self.mcp_client = None

            # 再次禁用stderr用于图形初始化
            sys.stderr = NullWriter()

            # 创建BaseGraph实例
            if self.mcp_client:
                self.graph = BaseGraph(
                    graph_name="mcp_graph",
                    nodes=[MCPNodes(mcp_client_manager=self.mcp_client)],
                    agent_template="agent_base",
                    parse_tool_args=True,
                    mcp_client=self.mcp_client,
                    verbose=False  # 关闭详细日志输出
                )
            else:
                # 如果没有MCP客户端，创建基础图
                self.graph = BaseGraph(
                    graph_name="basic_graph",
                    nodes=[],
                    agent_template="agent_base",
                    parse_tool_args=True,
                    verbose=False  # 关闭详细日志输出
                )

            # 恢复stderr
            sys.stderr = original_stderr

            # 生成线程ID
            import uuid
            self.thread_id = str(uuid.uuid4())

            self.console.print("[green]✅ AI 助手初始化完成！[/green]")
            self.console.print()

        except Exception as e:
            # 恢复stderr
            sys.stderr = original_stderr
            self.console.print(f"[red]❌ 初始化失败: {str(e)}[/red]")
            self.console.print("[yellow]请检查配置文件和环境设置[/yellow]")
            sys.exit(1)

    def print_help(self):
        """打印帮助信息"""
        self.console.print()
        self.console.print("[bold cyan]📖 帮助信息[/bold cyan]")
        self.console.print()
        self.console.print("[bold]可用命令：[/bold]")
        self.console.print("  [cyan]/help[/cyan]     - 显示此帮助信息")
        self.console.print("  [cyan]/quit[/cyan]     - 退出程序")
        self.console.print("  [cyan]/exit[/cyan]     - 退出程序")
        self.console.print("  [cyan]/clear[/cyan]    - 清屏")
        self.console.print("  [cyan]/new[/cyan]      - 开始新的对话会话")
        self.console.print("  [cyan]/status[/cyan]   - 显示当前状态")
        self.console.print("  [cyan]/history[/cyan]  - 显示对话历史")
        self.console.print()
        self.console.print("[bold]使用技巧：[/bold]")
        self.console.print("  1. 直接输入您的问题或需求，AI会分析并逐步完成")
        self.console.print("  2. 可以要求AI创建、修改文件，执行命令等")
        self.console.print("  3. AI会使用各种工具来完成复杂任务")
        self.console.print("  4. 对话历史会自动保存，新会话会包含之前的上下文")
        self.console.print()

    def print_status(self):
        """打印当前状态"""
        self.console.print()
        self.console.print("[bold green]📊 系统状态[/bold green]")
        self.console.print()
        self.console.print(f"当前目录: [dim]{os.getcwd()}[/dim]")
        self.console.print(f"线程ID: [dim]{self.thread_id}[/dim]")
        self.console.print(f"AI助手: {'[green]✅ 已连接[/green]' if self.graph else '[red]❌ 未连接[/red]'}")
        self.console.print(f"MCP客户端: {'[green]✅ 已启用[/green]' if self.mcp_client else '[red]❌ 未启用[/red]'}")
        self.console.print(f"对话历史: [dim]{len(conversation_cache.thread_order)} 个会话[/dim]")
        self.console.print()

    def print_history(self):
        """打印对话历史"""
        self.console.print()
        self.console.print("[bold cyan]📚 对话历史[/bold cyan]")
        self.console.print()

        if not conversation_cache.thread_order:
            self.console.print("[dim]暂无对话历史[/dim]")
            return

        for i, thread_id in enumerate(conversation_cache.thread_order, 1):
            messages = conversation_cache.get_history(thread_id)
            if messages:
                current_marker = " [bold green](当前)[/bold green]" if thread_id == self.thread_id else ""
                self.console.print(f"[bold]会话 {i}{current_marker}:[/bold] [dim]{thread_id}[/dim]")
                for msg in messages:  # 只显示最后2条消息
                    # 处理HumanMessage和AIMessage类型
                    content = msg.content
                    role = msg.type

                    role_color = "green" if role == "user" else "blue"
                    content_preview = content[:100] + "..." if len(content) > 100 else content
                    self.console.print(f"  [{role_color}]{role}[/{role_color}]: {content_preview}")
                self.console.print()

    def get_user_input(self, prompt_text: str) -> str:
        """获取用户输入，处理中文字符删除问题"""
        import readline

        # 设置readline的编码处理
        if hasattr(readline, 'set_startup_hook'):
            def startup_hook():
                readline.insert_text('')

            readline.set_startup_hook(startup_hook)

        try:
            # 使用input()而不是rich.Prompt来避免中文删除问题
            self.console.print(f"[bold green]{prompt_text}[/bold green]", end=" ")
            user_input = input()
            return user_input
        except (KeyboardInterrupt, EOFError):
            raise

    async def process_user_input(self, user_input: str) -> bool:
        """处理用户输入，返回是否继续运行"""

        # 处理命令
        if user_input.startswith('/'):
            command = user_input.lower().strip()

            if command in ['/quit', '/exit']:
                self.console.print("[yellow]👋 再见！感谢使用 Inscode AI Assistant[/yellow]")
                return False
            elif command == '/help':
                self.print_help()
                return True
            elif command == '/clear':
                self.console.clear()
                self.print_welcome()
                return True
            elif command == '/new':
                import uuid
                self.thread_id = str(uuid.uuid4())
                self.console.print("[green]✅ 已开始新的对话会话[/green]")
                return True
            elif command == '/status':
                self.print_status()
                return True
            elif command == '/history':
                self.print_history()
                return True
            else:
                self.console.print(f"[red]❌ 未知命令: {command}[/red]")
                self.console.print("[dim]输入 [cyan]/help[/cyan] 查看可用命令[/dim]")
                return True

        # 处理AI对话
        if not user_input.strip():
            return True

        try:
            # 显示思考状态
            self.console.print("[yellow]🤔 AI 正在思考中...[/yellow]")

            # 获取历史记录
            history = conversation_cache.get_all_history()

            # 构建请求
            request = GraphRequest(
                question=user_input,
                thread_id=self.thread_id,
                history=history,
                kwargs={
                    "cwd_dir": os.getcwd(),
                }
            )

            # 临时禁用stderr以避免处理过程中的日志输出
            sys.stderr = NullWriter()

            ai_response = ""  # 用于收集AI的完整回复

            # 使用流式处理
            async for message in self.graph.astream(request):
                if isinstance(message, List):
                    # 恢复stderr以显示任务完成信息
                    sys.stderr = original_stderr

                    # 获取并显示usage信息
                    messages = message
                    if messages:
                        all_usage_data = []
                        for ai_message in messages:
                            if hasattr(ai_message, 'response_metadata') and 'usage' in ai_message.response_metadata:
                                usage_data = ai_message.response_metadata['usage']
                                if isinstance(usage_data, list):
                                    all_usage_data.extend(usage_data)
                                elif usage_data:  # 单个usage对象
                                    all_usage_data.append(usage_data)

                        if all_usage_data:
                            for usage in all_usage_data:
                                model_name = usage.get('model', '未知模型')
                                input_tokens = usage.get('input_tokens', 0)
                                output_tokens = usage.get('output_tokens', 0)

                                self.console.print(
                                    f"[dim]📊 {model_name}: Input {input_tokens} Tokens | Output {output_tokens} Tokens[/dim]")

                        conversation_cache.add_message(self.thread_id, messages)

                    # 重新生成thread_id
                    import uuid
                    self.thread_id = str(uuid.uuid4())
                    break
                if not (isinstance(message, AgentStepLog) and message.meta["finish"] == True):
                    continue
                output = message.output
                if message.type == 'tool':
                    try:
                        # 解析工具执行内容
                        if not isinstance(output, GraphNodeParam):
                            continue

                        if output.type == 'use_mcp_tool':
                            params = output.params
                            tool_name = params.get('tool_name', '未知工具')
                            server_name = params.get('server_name', '未知服务')

                            # 恢复stderr以显示工具执行信息
                            sys.stderr = original_stderr
                            # 显示工具执行信息
                            self.console.print(
                                f"[yellow]🔧 正在执行工具: [bold cyan]{tool_name}[/bold cyan] (服务: {server_name})[/yellow]")
                            # 再次禁用stderr
                            sys.stderr = NullWriter()
                    except (json.JSONDecodeError, AttributeError):
                        pass
                else:
                    # 恢复stderr以显示AI回复
                    sys.stderr = original_stderr
                    self.console.print()
                    # AI回复在右侧显示
                    self.console.print(Panel(
                        Markdown(output),
                        title="[bold blue]🤖 Inscode AI[/bold blue]",
                        title_align="left",
                        border_style="blue",
                        padding=(1, 2)
                    ))

                    # 收集AI回复内容
                    ai_response = output

                    # 再次禁用stderr
                    sys.stderr = NullWriter()

            # 恢复stderr
            sys.stderr = original_stderr

            # 显示会话结束分隔符
            self.console.print()
            self.console.print(Rule(title="[dim]会话结束[/dim]", style="dim"))
            self.console.print()

        except Exception as e:
            # 恢复stderr
            sys.stderr = original_stderr
            self.console.print(f"[red]❌ 处理请求时出错: {str(e)}[/red]")

        return True

    async def run_interactive_session(self):
        """运行交互式会话"""
        self.console.print()

        while True:
            try:
                # 使用自定义输入方法处理中文字符删除问题
                user_input = self.get_user_input("👤 You")

                # 处理输入
                should_continue = await self.process_user_input(user_input)
                if not should_continue:
                    break

            except KeyboardInterrupt:
                self.console.print("\n[yellow]👋 再见！感谢使用 Inscode AI Assistant[/yellow]")
                break
            except EOFError:
                self.console.print("\n[yellow]👋 再见！感谢使用 Inscode AI Assistant[/yellow]")
                break

    async def cleanup(self):
        """清理资源"""
        try:
            # 临时禁用stderr以避免清理过程中的日志
            sys.stderr = NullWriter()
            if self.mcp_client:
                await self.mcp_client.cleanup()
            # 恢复stderr
            sys.stderr = original_stderr
        except Exception as e:
            # 恢复stderr
            sys.stderr = original_stderr
            pass  # 静默处理清理错误

    async def run(self):
        """运行客户端"""
        try:
            # 检查必需文件
            if not self.check_required_files():
                sys.exit(1)

            # 加载环境变量
            if not self.load_environment():
                sys.exit(1)

            self.print_welcome()
            await self.initialize_graph()
            await self.run_interactive_session()
        finally:
            await self.cleanup()


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Inscode AI Assistant 命令行客户端")
    parser.add_argument(
        "--version",
        action="version",
        version="Inscode AI Assistant v1.0.0"
    )

    args = parser.parse_args()

    # 创建并运行客户端
    client = InsCodeClient()
    await client.run()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n👋 程序已退出")
    except Exception as e:
        print(f"❌ 程序运行出错: {e}")
        sys.exit(1)
